﻿using System;
using System.Linq;
using Bettzueche.RLGlue.RLComponent;
using Bettzueche.RLGlue.RLTypes;
using Bettzueche.RLGlue.TaskSpec;
using System.Diagnostics;
using System.IO;
using Bettzueche.RLGlue.Codec.Util;

namespace RunMinesSarsaExperiment {
    /// <summary>
    /// This is a very simple Sarsa agent for discrete-action, discrete-state
    /// environments.  It uses epsilon-greedy exploration.<para>
    /// We've made a decision to store the previous action and observation in
    /// their raw form, as structures.  This code could be simplified and you
    /// could store them just as ints.</para><para>
    /// @author Brian Tanner, Copyright 2009</para>
    /// </summary>
    class SampleSarsaAgent : IAgent {

        #region Fields

        Random randGenerator = new Random();
        GeneralAction lastAction;
        GeneralState lastObservation;
        double[,] valueFunction = null;
        double sarsa_stepsize = 0.1;
        double sarsa_epsilon = 0.1;
        double sarsa_gamma = 1.0;
        int numActions = 0;
        int numStates = 0;
        bool policyFrozen = false;
        bool exploringFrozen = false;

        #endregion

        /// <summary>
        /// "This is a trick we can use to make the agent easily loadable.  Using this
        /// trick you can directly execute the class and it will load itself through
        /// AgentLoader and connect to the rl_glue server."
        /// </summary>
        /// <param name="args"></param>
        static void Main(string[] args) {
            AgentLoader theLoader = AgentLoader.LoadAgent(new SampleSarsaAgent());
            theLoader.Run();
        }

        #region IAgent Members

        /// <summary>
        /// "Parse the task spec, make sure it is only 1 integer observation and
        /// action, and then allocate the valueFunction." (BTanner)
        /// </summary>
        /// <param name="task_spec"></param>
        public void InitAgent(string task_spec) {
            ITaskSpec theTaskSpec = TaskSpecFactory.GetTaskSpec(task_spec);

            /* Lots of assertions to make sure that we can handle this problem.  */
            Debug.Assert(theTaskSpec.Observations.Discrete.Count() == 1);
            Debug.Assert(theTaskSpec.Observations.Continuous.Count() == 0);
            //assert (!theTaskSpec.getDiscreteObservationRange(0).hasSpecialMinStatus());
            //assert (!theTaskSpec.getDiscreteObservationRange(0).hasSpecialMaxStatus());
            numStates = theTaskSpec.Observations.Discrete.ElementAt(0).Max + 1;

            Debug.Assert(theTaskSpec.Actions.Discrete.Count() == 1);
            Debug.Assert(theTaskSpec.Actions.Continuous.Count() == 0);
            //assert (!theTaskSpec.getDiscreteActionRange(0).hasSpecialMinStatus());
            //assert (!theTaskSpec.getDiscreteActionRange(0).hasSpecialMaxStatus());
            numActions = theTaskSpec.Actions.Discrete.ElementAt(0).Max + 1;

            sarsa_gamma=theTaskSpec.Discountfactor;

            valueFunction = new double[numActions, numStates];
        }

        /// <summary>
        /// "Choose an action e-greedily from the value function and store the action and observation." (BTanner)
        /// </summary>
        /// <param name="observation"></param>
        /// <returns></returns>
        public GeneralAction StartAgent(GeneralState observation) {
            int newActionInt = egreedy(observation.IntValues[0]);

            /*
             * Create a structure to hold 1 integer action
             * and set the value
             */
            GeneralAction returnAction = new GeneralAction(new int[]{newActionInt});

            lastAction = new GeneralAction(returnAction);
            lastObservation = new GeneralState(observation);

            return returnAction;
        }

        /// <summary>
        /// "Choose an action e-greedily from the value function and store the action
        /// and observation.  Update the valueFunction entry for the last
        /// state,action pair." (BTanner)
        /// </summary>
        /// <param name="reward"></param>
        /// <param name="observation"></param>
        /// <returns></returns>
        public GeneralAction StepAgent(double reward, GeneralState observation) {
            int newStateInt = observation.IntValues[0];
            int lastStateInt = lastObservation.IntValues[0];
            int lastActionInt = lastAction.IntValues[0];

            int newActionInt = egreedy(newStateInt);

            double Q_sa = valueFunction[lastActionInt,lastStateInt];
            double Q_sprime_aprime = valueFunction[newActionInt,newStateInt];

            double new_Q_sa = Q_sa + sarsa_stepsize * (reward + sarsa_gamma * Q_sprime_aprime - Q_sa);
            /*	Only update the value function if the policy is not frozen */
            if (!policyFrozen) {
                valueFunction[lastActionInt,lastStateInt] = new_Q_sa;
            }

            /* Creating the action a different way to showcase variety */
            GeneralAction returnAction = new GeneralAction();
            returnAction.IntValues = new int[] { newActionInt };

            lastAction = new GeneralAction(returnAction);
            lastObservation = new GeneralState(observation);

            return returnAction;
        }

        /// <summary>
        /// "The episode is over, learn from the last reward that was received." (BTanner)
        /// </summary>
        /// <param name="reward"></param>
        public void EndAgent(double reward) {
            int lastStateInt = lastObservation.IntValues[0];
            int lastActionInt = lastAction.IntValues[0];

            double Q_sa = valueFunction[lastActionInt,lastStateInt];
            double new_Q_sa = Q_sa + sarsa_stepsize * (reward - Q_sa);

            /*	Only update the value function if the policy is not frozen */
            if (!policyFrozen) {
                valueFunction[lastActionInt,lastStateInt] = new_Q_sa;
            }
            lastObservation = null;
            lastAction = null;
        }

        /// <summary>
        /// "Release memory that is no longer required/used." (BTanner)
        /// </summary>
        public void CleanupAgent() {
            lastAction = null;
            lastObservation = null;
            valueFunction = null;
        }

        /// <summary>
        /// "This agent responds to some simple messages for freezing learning and
        /// saving/loading the value function to a file." (BTanner)
        /// </summary>
        /// <param name="message"></param>
        /// <returns></returns>
        public string MessageAgent(string message) {
            if (message == "freeze learning") {
                policyFrozen = true;
                return "message understood, policy frozen";
            }
            if (message == "unfreeze learning") {
                policyFrozen = false;
                return "message understood, policy unfrozen";
            }
            if (message == "freeze exploring") {
                exploringFrozen = true;
                return "message understood, exploring frozen";
            }
            if (message == "unfreeze exploring") {
                exploringFrozen = false;
                return "message understood, exploring unfrozen";
            }
            if (message.StartsWith("save_policy")) {
                string[] parts = message.Split(' ');
                saveValueFunction(parts[1]);
                Console.WriteLine("Saved.");
                return "message understood, saving policy";
            }
            if (message.StartsWith("load_policy")) {
                string[] parts = message.Split(' ');
                loadValueFunction(parts[1]);
                Console.WriteLine("Loaded.");
                return "message understood, loading policy";
            }

            return "SampleSarsaAgent(Java) does not understand your message.";
        }

        #endregion

        #region Private Methods

        /// <summary>
        /// "Selects a random action with probability 1-sarsa_epsilon,
        /// and the action with the highest value otherwise.  This is a
        /// quick'n'dirty implementation, it does not do tie-breaking." (BTanner)
        /// </summary>
        /// <param name="theState"></param>
        /// <returns></returns>
        private int egreedy(int theState) {
            if (!exploringFrozen) {
                if (randGenerator.NextDouble() <= sarsa_epsilon) {
                    return randGenerator.Next(numActions);
                }
            }

            /*otherwise choose the greedy action*/
            int maxIndex = 0;
            for (int a = 1; a < numActions; a++) {
                if (valueFunction[a,theState] > valueFunction[maxIndex,theState]) {
                    maxIndex = a;
                }
            }
            return maxIndex;
        }


        /// <summary>
        /// "Dumps the value function to a file named theFileName.  Not fancy. Must be
        /// called after init but before cleanup." (BTanner)
        /// </summary>
        /// <param name="theFileName"></param>
        void saveValueFunction(string theFileName) {
            using (BinaryWriter writer = new BinaryWriter(File.Open(theFileName,FileMode.Create))) {
                for (int a = 0; a < numActions; a++) {
                    for (int s = 0; s < numStates; s++) {
                        try {
                            writer.Write(valueFunction[a, s]);
                        }
                        catch (ObjectDisposedException) {
                            Console.Error.WriteLine("Unexpected dispose of filestream while saving value function.");
                        }
                        catch (IOException ex) {
                            Console.Error.WriteLine("Problem writing value function to file:: " + ex.Message);
                        }
                    }
                }
            }
        }

        /// <summary>
        /// "Loads the value function from a file named theFileName.  Must be called after
        /// init but before cleanup." (BTanner)
        /// </summary>
        /// <param name="theFileName"></param>
        void loadValueFunction(string theFileName) {
            using (BinaryReader reader = new BinaryReader(File.Open(theFileName,FileMode.OpenOrCreate))) {
                for (int a = 0; a < numActions; a++) {
                    for (int s = 0; s < numStates; s++) {
                        
                        try {	        
                            valueFunction[a,s] = reader.ReadDouble();
                        }
                        catch (EndOfStreamException) {
                            Console.Error.WriteLine("Unexpected end of file, while reading value function.");
                        }
                        catch (ObjectDisposedException) {
                            Console.Error.WriteLine("Unexpected dispose of filestream while reading value function.");
                        }
                        catch (IOException ex) {
                            Console.Error.WriteLine("Problem reading value function from file:: " + ex.Message);
                        }
                    }
                }
            }
        }

        #endregion
    }
}
